#include "compi/compi.h"

#include <gtest/gtest.h>

#include <chrono>
#include <string>
#include <thread>

using namespace compi;

// 前置声明：辅助函数在文件底部定义
static bool ns_consistent(MPI_Comm comm, TagNamespace ns);
static bool wait_with_timeout(TransportRequest & req, int timeout_ms);

// 全局测试环境，由自定义 main 管理 MPI 生命周期
static Environment * g_test_env = nullptr;

class CancelPollGroupTest : public ::testing::Test
{
protected:
    void SetUp() override
    {
        env_ = g_test_env;
        // 使用 Context::for_comm 获取非拥有的上下文实例
        comm_ = &Context::for_comm(MPI_COMM_WORLD, *env_);
    }
    void TearDown() override
    {
        // 通过 Context 释放与 communicator 关联的资源
        Context::destroy_for_comm(MPI_COMM_WORLD);
        // 非拥有指针清空，保留全局 Environment 直至进程退出
        comm_ = nullptr;
        env_ = nullptr;
    }
    Environment * env_{};
    Context * comm_{};
};

// 自定义 main：统一初始化与析构，避免 mpiexec 对 Finalize 的误判
int
main(int argc, char ** argv)
{
    ::testing::InitGoogleTest(&argc, argv);
    g_test_env = new Environment();
    const int code = RUN_ALL_TESTS();
    delete g_test_env;
    g_test_env = nullptr;
    return code;
}

// 用例1：在未匹配发布阶段取消（固定类型）
TEST_F(CancelPollGroupTest, CancelBeforeMatchFixed)
{
    auto & alloc = comm_->tag_allocator();
    auto ns = alloc.create_namespace("cancel_poll_ns_fixed");

    // 命名空间需跨 rank 一致
    if (!ns_consistent(comm_->comm(), ns))
        GTEST_SKIP() << "TagNamespace IDs are inconsistent across ranks; skipping to avoid hang.";

    const int rank = comm_->rank();
    const int size = comm_->size();
    const int src = (rank - 1 + size) % size;

    int recv_val = -1;
    TransportRequest r = irecv(*comm_, src, recv_val, ns);

    // 收集处理器上下文（应在取消完成时被调用，cancelled=true）
    RequestStatus captured{};
    bool processor_called = false;
    r.set_strategy(StrategyUnpack | StrategyProcess | StrategyRelease);
    r.set_processor(
        [&](void * payload, const RequestStatus & st)
        {
            (void)payload;
            captured = st;
            processor_called = true;
        });

    // 不发送任何消息，直接取消
    ASSERT_TRUE(r.cancel());
    // 取消完成后应可立即查询状态
    auto st = r.status();
    EXPECT_TRUE(st.cancelled);
    EXPECT_FALSE(st.error.has_value());
    EXPECT_EQ(st.tag, MPI_ANY_TAG);
    EXPECT_TRUE(processor_called);
    EXPECT_TRUE(captured.cancelled);
}

// 用例2：poll 推进并完成（固定类型，带处理器上下文）
TEST_F(CancelPollGroupTest, PollProgressFixedWithProcessor)
{
    auto & alloc = comm_->tag_allocator();
    auto ns = alloc.create_namespace("poll_progress_ns_fixed");

    if (!ns_consistent(comm_->comm(), ns))
        GTEST_SKIP() << "TagNamespace IDs are inconsistent across ranks; skipping to avoid hang.";

    const int rank = comm_->rank();
    const int size = comm_->size();
    const int dest = (rank + 1) % size;
    const int src = (rank - 1 + size) % size;

    int recv_val = -1;
    TransportRequest r = irecv(*comm_, src, recv_val, ns);

    // 发送前同步，确保接收已发布
    MPI_Barrier(comm_->comm());

    int payload = rank * 100 + 7;
    auto sreq = isend(*comm_, dest, payload, ns, SendMode::Default);

    // 处理器：记录状态与字节数
    bool processor_called = false;
    RequestStatus captured{};
    r.set_strategy(StrategyUnpack | StrategyProcess | StrategyRelease);
    r.set_processor(
        [&](void * payload_ptr, const RequestStatus & st)
        {
            (void)payload_ptr;
            captured = st;
            processor_called = true;
        });

    // 使用 poll 推进直到完成
    ASSERT_TRUE(wait_with_timeout(r, 10000)) << "Receive did not complete in time";
    ASSERT_TRUE(sreq.wait().is_success());

    EXPECT_TRUE(processor_called);
    EXPECT_FALSE(captured.cancelled);
    EXPECT_FALSE(captured.error.has_value());
    // 对固定类型，count_bytes 应为 sizeof(int)
    EXPECT_EQ(captured.count_bytes, static_cast<int>(sizeof(int)));
}

// 用例3：RequestGroup 组合：一个完成，一个取消
TEST_F(CancelPollGroupTest, RequestGroupMixedCompleteAndCancel)
{
    auto & alloc = comm_->tag_allocator();
    auto nsA = alloc.create_namespace("group_ns_A_fixed");
    auto nsB = alloc.create_namespace("group_ns_B_fixed");

    if (!ns_consistent(comm_->comm(), nsA) || !ns_consistent(comm_->comm(), nsB))
        GTEST_SKIP() << "TagNamespace IDs are inconsistent across ranks; skipping to avoid hang.";

    const int rank = comm_->rank();
    const int size = comm_->size();
    const int dest = (rank + 1) % size;
    const int src = (rank - 1 + size) % size;

    int recvA = -1, recvB = -1;
    TransportRequest rA = irecv(*comm_, src, recvA, nsA);
    TransportRequest rB = irecv(*comm_, src, recvB, nsB);
    rA.set_strategy(StrategyUnpack | StrategyProcess | StrategyRelease);
    rB.set_strategy(StrategyUnpack | StrategyProcess | StrategyRelease);

    bool procA = false, procB = false;
    RequestStatus stA{}, stB{};
    rA.set_processor(
        [&](void * payload, const RequestStatus & st)
        {
            (void)payload;
            stA = st;
            procA = true;
        });
    rB.set_processor(
        [&](void * payload, const RequestStatus & st)
        {
            (void)payload;
            stB = st;
            procB = true;
        });

    // 同步后仅发送 B
    MPI_Barrier(comm_->comm());
    int payloadB = rank * 1000 + 5;
    auto sreqB = isend(*comm_, dest, payloadB, nsB, SendMode::Default);

    RequestGroup group;
    group.add(rA);
    group.add(rB);

    // 推进所有请求，直到 B 完成
    auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(10000);
    bool b_done = false;
    while (std::chrono::steady_clock::now() < deadline)
    {
        group.poll_all();
        auto done = group.test_some();
        for (auto idx : done)
        {
            if (idx == 1) // rB 完成
                b_done = true;
        }
        if (b_done)
            break;
        std::this_thread::sleep_for(std::chrono::milliseconds(2));
    }
    ASSERT_TRUE(b_done) << "Request B did not complete in time";

    // 取消未完成的 A
    ASSERT_TRUE(rA.cancel());

    // 组内等待全部完成（B 已完成，A 已取消/完成）
    ASSERT_TRUE(group.wait_all());
    ASSERT_TRUE(sreqB.wait().is_success());

    // 验证处理器调用与状态
    EXPECT_TRUE(procB);
    EXPECT_FALSE(stB.cancelled);
    EXPECT_FALSE(stB.error.has_value());
    EXPECT_TRUE(procA);
    EXPECT_TRUE(stA.cancelled);
    EXPECT_FALSE(stA.error.has_value());
}

// 工具函数：检查 TagNamespace 的 ID 在所有 rank 间是否一致
static bool
ns_consistent(MPI_Comm comm, TagNamespace ns)
{
    int world_size = 0;
    MPI_Comm_size(comm, &world_size);
    int my_id = static_cast<int>(ns.id());
    std::vector<int> all_ids(world_size, 0);
    MPI_Allgather(&my_id, 1, MPI_INT, all_ids.data(), 1, MPI_INT, comm);
    for (int i = 1; i < world_size; ++i)
        if (all_ids[i] != all_ids[0])
            return false;
    return true;
}

// 等待工具：为请求添加超时，避免测试在异常情况下卡住
static bool
wait_with_timeout(TransportRequest & req, int timeout_ms)
{
    using namespace std::chrono;
    const auto deadline = steady_clock::now() + milliseconds(timeout_ms);
    while (true)
    {
        if (req.completed())
            return req.wait().is_success();
        if (steady_clock::now() > deadline)
            return false;
        std::this_thread::sleep_for(milliseconds(2));
    }
}